package com.aiwiown.snackmq.storage;

import com.aiwiown.snackmq.common.message.Message;
import com.aiwiown.snackmq.storage.segment.LogSegment;
import com.aiwiown.snackmq.storage.segment.SegmentFullException;
import lombok.extern.slf4j.Slf4j;

import java.io.File;
import java.io.IOException;
import java.util.Arrays;
import java.util.Map;
import java.util.concurrent.ConcurrentNavigableMap;
import java.util.concurrent.ConcurrentSkipListMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * 管理单个分区的日志，包括其所有的日志段 (LogSegment)。
 */
@Slf4j
public class PartitionLog {

    private final String partitionPath;
    private final long segmentSize;
    private final int indexInterval;
    // 【新增】存储恢复策略配置
    private final boolean forceFullLogRecovery;

    private final ConcurrentNavigableMap<Long, LogSegment> segments = new ConcurrentSkipListMap<>();
    private final AtomicLong nextOffset;
    private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();

    public PartitionLog(String partitionPath, long segmentSize, int indexInterval, boolean forceFullLogRecovery) {
        this.partitionPath = partitionPath;
        this.segmentSize = segmentSize;
        this.indexInterval = indexInterval;
        this.forceFullLogRecovery = forceFullLogRecovery; // 【新增】
        this.nextOffset = new AtomicLong(0);
    }

    public void load() throws IOException {
        lock.writeLock().lock();
        try {
            File dir = new File(partitionPath);
            if (!dir.exists() && !dir.mkdirs()) {
                throw new IOException("Failed to create partition directory: " + partitionPath);
            }

            File[] logFiles = dir.listFiles((d, name) -> name.endsWith(".log"));
            if (logFiles == null || logFiles.length == 0) {
                createNewSegment(0);
                this.nextOffset.set(0);
                return;
            }

            Arrays.sort(logFiles, (f1, f2) -> {
                long off1 = Long.parseLong(f1.getName().replace(".log", ""));
                long off2 = Long.parseLong(f2.getName().replace(".log", ""));
                return Long.compare(off1, off2);
            });

            for (File logFile : logFiles) {
                try {
                    long baseOffset = Long.parseLong(logFile.getName().replace(".log", ""));
                    LogSegment segment = new LogSegment(partitionPath, baseOffset, segmentSize, indexInterval, this.forceFullLogRecovery);
                    segments.put(baseOffset, segment);
                } catch (NumberFormatException e) {
                    log.warn("Skipping file with invalid name in partition directory {}: {}", partitionPath, logFile.getName());
                }
            }

            LogSegment lastSegment = segments.lastEntry().getValue();
            this.nextOffset.set(lastSegment.getBaseOffset() + lastSegment.getEntryCount().get());
            log.info("Successfully loaded partition log from {}. Next offset is {}.", partitionPath, nextOffset.get());

        } finally {
            lock.writeLock().unlock();
        }
    }

    public long append(Message message) throws IOException {
        lock.writeLock().lock();
        try {
            long logicalOffset = nextOffset.get();
            LogSegment segment = latestSegment();
            try {
                segment.append(message, logicalOffset);
            } catch (SegmentFullException e) {
                LogSegment newSegment = createNewSegment(logicalOffset);
                newSegment.append(message, logicalOffset);
            }
            return nextOffset.getAndIncrement();
        } finally {
            lock.writeLock().unlock();
        }
    }

    public Message pull(long offset) throws IOException {
        lock.readLock().lock();
        try {
            Map.Entry<Long, LogSegment> entry = segments.floorEntry(offset);
            if (entry == null) {
                return null;
            }
            LogSegment segment = entry.getValue();
            return segment.read(offset);
        } finally {
            lock.readLock().unlock();
        }
    }

    /**
     * 关闭此分区日志并释放其所有日志段的资源。
     * 这是解决编译错误的关键方法。
     */
    public void close() throws IOException {
        lock.writeLock().lock();
        try {
            for (LogSegment segment : segments.values()) {
                segment.close();
            }
            segments.clear();
            log.info("Partition log at {} closed.", partitionPath);
        } finally {
            lock.writeLock().unlock();
        }
    }

    private LogSegment latestSegment() {
        return segments.lastEntry().getValue();
    }

    private LogSegment createNewSegment(long baseOffset) throws IOException {
        LogSegment newSegment = new LogSegment(partitionPath, baseOffset, segmentSize, indexInterval, this.forceFullLogRecovery);
        segments.put(baseOffset, newSegment);
        log.info("Created new segment with base offset {} for partition {}", baseOffset, partitionPath);
        return newSegment;
    }
}